#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : chat_model 
# @File    : ernie.py
# @IDE     : PyCharm 
# @Author  :ZH
# @Time    : 2025/1/8 13:03
import asyncio
import json
import time
from typing import List

from apps.model_utils.base import BaseModelUtil
from apps.utils import async_http_client


class Ernie(BaseModelUtil):

    def __init__(
            self, api_key, prompt: str = None, message: List[dict] = None,
            temperature: float = 0.95, top_p: float = 0.7, max_tokens=4096,
            penalty_score: float = 1.0, *args, **kwargs):
        self.api_key = api_key
        self.prompt = prompt
        self.message = message if message else []
        self.temperature = temperature
        self.top_p = top_p
        self.max_tokens = max_tokens
        self.penalty_score = penalty_score
        self.usage_tokens = 0

    async def chat(self, query, base_url, model_name, model_object, **kwargs):
        access_token = await self.get_access_token(model_object=model_object)
        base_url = base_url + f'/{model_name}' if base_url[-1] != '/' else base_url + f'{model_name}'
        params = {
            'access_token': access_token,
        }
        message = self.get_message()
        message.append({"role": "user", "content": query})
        response = await async_http_client.post(
            url=base_url, params=params, json={
                'messages': message,
                'temperature': self.temperature,
                'top_p': self.top_p,
                # 'max_output_tokens': self.max_tokens,
                'system': self.prompt,
                'penalty_score': self.penalty_score
            },
            headers=self._header()
        )
        print(response)
        self.usage_tokens = response.get("usage", {}).get('total_tokens')
        return response.get('result')

    async def get_access_token(self, model_object):
        self.api_key = json.loads(self.api_key)
        # 判断上次生成的时间戳距离现在的时间间隔
        if self.api_key.get('generation_time') and self.api_key.get('access_token'):
            if int(time.time() * 1000) - self.api_key['generation_time'] < 15 * 24 * 60 * 60:
                print('当前未过期')
                return self.api_key['access_token']

        # 获取access_token
        client_id = self.api_key.get('client_id')  # AK
        client_secret = self.api_key.get('client_secret')  # SK
        url = f"https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={client_id}&client_secret={client_secret}"
        headers = {
            "Content-Type": "application/json",
        }
        response = await async_http_client.post(url, headers=headers)
        access_token = response.get('access_token')
        self.api_key['access_token'] = access_token
        self.api_key['generation_time'] = int(time.time() * 1000)
        model_object.api_key = json.dumps(self.api_key, ensure_ascii=False)
        await model_object.save()
        return access_token

    def _header(self) -> dict:
        return {"Content-Type": "application/json"}

    def get_message(self):
        if self.message:
            _ = []
            for i in self.message:
                _.append({"role": "user", "content": i.get('query')})
                _.append({"role": "assistant", "content": i.get('answer')})
            return _
        else:
            return []

    async def get_usage_tokens(self):
        return self.usage_tokens

